from collections import namedtuple
import cv2


# Set up text parameters
OWL_FONT = cv2.FONT_HERSHEY_SIMPLEX


# https://claude.ai/chat/a9c2c935-26b6-4102-ac9a-5ac95dd7d8a0
def draw_bbox_with_label(image, bbox, label, color=(255, 255, 0), thickness=1, font_scale=0.7):
    """
    Draw a bounding box with a text label on an image.

    Parameters:
    -----------
    image : numpy.ndarray
        The input image on which to draw
    bbox : tuple
        Bounding box coordinates as (x_min, y_min, x_max, y_max)
    label : str
        Text label to display
    color : tuple, optional
        BGR color tuple for box and text, default is green (0, 255, 0)
    thickness : int, optional
        Line thickness for the bounding box, default is 2
    font_scale : float, optional
        Font scale for the text label, default is 0.7

    Returns:
    --------
    numpy.ndarray
        Image with bounding box and label drawn
    """

    # Make a copy of the image to avoid modifying the original
    img_with_bbox = image#image.copy()

    # Extract coordinates
    x_min, y_min, x_max, y_max = [int(b) for b in bbox]

    width, height = image.shape[2], image.shape[1]
    assert x_min in range(0, width), f"Bounding box x_min is out of the image ({x_min})"
    assert y_min in range(0, height), f"Bounding box y_min is out of the image ({y_min})"
    assert x_max in range(0, width), f"Bounding box x_max is out of the image ({x_max})"
    assert y_max in range(0, height), f"Bounding box y_max is out of the image ({y_max})"

    # Draw the bounding box
    cv2.rectangle(img_with_bbox, (x_min, y_min), (x_max, y_max), color, thickness)

    # Get text size
    (text_width, text_height), baseline = cv2.getTextSize(label, OWL_FONT, font_scale, thickness)

    # Draw background rectangle for the text
    cv2.rectangle(img_with_bbox,
                  (x_min, y_min - text_height - 10),
                  (x_min + text_width + 10, y_min),
                  color, -1)  # -1 thickness fills the rectangle

    # Draw the text
    cv2.putText(img_with_bbox,
                label,
                (x_min + 5, y_min - 5),
                OWL_FONT, font_scale,
                (255, 255, 255),  # White text
                thickness,
                cv2.LINE_AA)

    return img_with_bbox


# NOTE: Looks like the import for Owlv2ObjectDetectionOutput is hidden, but thanks to duck typing
# we can mock the necessary variables
Owlv2ObjectDetectionOutput = namedtuple('Owlv2ObjectDetectionOutput', ['logits', 'pred_boxes', 'target_pred_boxes'])
